"""
    StateVectors

Module implementing state vector operations for GSI variational data assimilation.
This module is ported from the GSI Fortran `state_vectors.f90`, `control2state.f90`,
and `control2model.f90` modules and provides efficient implementations of state vector
operations, transformations, and memory management.

The state vector represents the complete atmospheric state in physical space:
```
x = [u, v, T, q, ps, oz, cw, ...]ᵀ
```

Where variables include:
- `u,v`: Horizontal wind components
- `T`: Temperature  
- `q`: Specific humidity
- `ps`: Surface pressure
- `oz`: Ozone mixing ratio
- `cw`: Cloud water mixing ratio

This module handles the critical transformations between control space and physical space:

1. **Control-to-State**: Transform control variables to physical state variables
   - Streamfunction/velocity potential → wind components (u,v)
   - Normalized variables → physical variables  
   - Apply background error covariance transformations

2. **State-to-Model**: Prepare state variables for model integration
   - Apply balance constraints (geostrophic, hydrostatic)
   - Ensure physical consistency (thermodynamic relations)
   - Handle cloud microphysics transformations

# Mathematical Framework

The state vector transformation follows:
```
x_state = B^(1/2) * x_control + x_background
```

Where:
- `x_state`: Physical state variables
- `x_control`: Control space variables (optimized during minimization)
- `B^(1/2)`: Square root of background error covariance matrix
- `x_background`: Background (first guess) state

# Core Operations

## State Vector Management
- Allocation and deallocation with memory tracking
- Norm computation (min, max, mean) across all variables
- Dot product operations for gradient calculations
- Random initialization for ensemble methods

## Variable Transformations
- Streamfunction/velocity potential to u,v wind components
- Relative humidity normalization and conversion to specific humidity
- Virtual temperature to sensible temperature conversion
- Cloud water to hydrometeor species transformation
- Pressure coordinate transformations

## Grid Support  
- Global and regional grid configurations
- Multi-level vertical coordinates
- MPI domain decomposition support
- Efficient memory layout for parallel processing

# Usage Example

```julia
using GSICoreAnalysis.StateVectors

# Configure analysis system
config = AnalysisConfig(
    grid_size = (721, 1440, 64),    # lat, lon, levels
    ensemble_size = 40,
    use_hybrid = true
)

# Create state vector
state = StateVector(config)

# Transform control vector to state
control_to_state!(control_vec, state, config)

# Apply model transformations
# state_to_model!(state, model_state, config)
#
# # Compute norms
# norms = compute_state_norms(state)
# println("Temperature norm: ", norms["t"])
#
# # Dot product for gradient calculations
# gradient_contrib = dot_product(state1, state2, variable="t")
```
"""
module StateVectors

using LinearAlgebra
using SparseArrays
using StaticArrays
using Distributed
using FFTW
using MPI
using Printf

using ..GSICoreAnalysis: AbstractAnalysisConfig, AbstractStateVector, AnalysisVariables, default_analysis_variables, GridDimensions

export StateVector
export allocate_state!, deallocate_state!
export control_to_state!, state_to_model!  
export compute_state_norms, dot_product_state, set_random_state!
export inquire_state_memory, print_state_norms
export streamfunction_to_winds!, pressure_transforms!, humidity_transforms!
export balance_constraints!, cloud_microphysics_transforms!

"""
    StateVector{T<:AbstractFloat}

State vector for GSI variational data assimilation containing physical variables.

# Fields

## Core Physical Variables (3D)
- `u::Array{T,3}`: Zonal wind component [m/s]
- `v::Array{T,3}`: Meridional wind component [m/s]  
- `t::Array{T,3}`: Temperature [K]
- `q::Array{T,3}`: Specific humidity [kg/kg]
- `prse::Array{T,3}`: Pressure on edges [Pa]
- `oz::Array{T,3}`: Ozone mixing ratio [kg/kg]

## Cloud Variables (3D)
- `cw::Array{T,3}`: Cloud water mixing ratio [kg/kg]
- `ql::Array{T,3}`: Liquid water content [kg/kg]
- `qi::Array{T,3}`: Ice water content [kg/kg]
- `qr::Array{T,3}`: Rain water content [kg/kg]
- `qs::Array{T,3}`: Snow content [kg/kg]
- `qg::Array{T,3}`: Graupel content [kg/kg]

## Surface Variables (2D)
- `ps::Matrix{T}`: Surface pressure [Pa]
- `sst::Matrix{T}`: Sea surface temperature [K]
- `stl::Matrix{T}`: Skin temperature over land [K]
- `sti::Matrix{T}`: Skin temperature over ice [K]

## Diagnostic/Derived Variables (3D)
- `tv::Array{T,3}`: Virtual temperature [K]
- `tsen::Array{T,3}`: Sensible temperature [K]
- `w::Array{T,3}`: Vertical velocity [m/s]

## Auxiliary Variables (2D)
- `gust::Matrix{T}`: Surface wind gust [m/s]
- `vis::Matrix{T}`: Visibility [m]
- `pblh::Matrix{T}`: Planetary boundary layer height [m]
- `wspd10m::Matrix{T}`: 10-meter wind speed [m/s]
- `tcamt::Matrix{T}`: Total cloud amount [fraction]
- `lcbas::Matrix{T}`: Low cloud base [m]

## Chemical Species (3D) - Optional
- `co::Array{T,3}`: Carbon monoxide mixing ratio [kg/kg]
- `co2::Array{T,3}`: Carbon dioxide mixing ratio [kg/kg]

## Configuration and Metadata
- `config::AbstractAnalysisConfig`: Analysis configuration
- `variables::AnalysisVariables`: Variable metadata and indexing
- `grid::GridDimensions`: Grid dimension information
- `is_allocated::Bool`: Memory allocation status
- `memory_usage::Dict{String,Int}`: Memory usage tracking by variable

## MPI Information
- `mpi_rank::Int`: MPI process rank
- `mpi_size::Int`: Total number of MPI processes
- `local_grid_start::Tuple{Int,Int,Int}`: Starting indices for local domain
- `local_grid_end::Tuple{Int,Int,Int}`: Ending indices for local domain

# Mathematical Properties

The state vector represents the atmospheric state at analysis time:
```
x = [u(i,j,k), v(i,j,k), T(i,j,k), q(i,j,k), ps(i,j), ...]ᵀ
```

Where indices represent:
- `i`: Longitude index (1:nx)
- `j`: Latitude index (1:ny)  
- `k`: Vertical level index (1:nz)

Total state vector dimension: N = nx*ny*(nz*n3d + n2d)
Where n3d is number of 3D variables and n2d is number of 2D variables.

# Physical Constraints

State variables must satisfy:
1. **Hydrostatic Balance**: ∂p/∂z = -ρg
2. **Thermodynamic Relations**: pV = nRT
3. **Continuity**: ∇·(ρV) + ∂ρ/∂t = 0
4. **Moisture Conservation**: 0 ≤ q ≤ q_sat(T,p)
5. **Positive Definiteness**: ps > 0, T > 0, mixing ratios ≥ 0

# Memory Layout

Variables are stored in column-major order optimized for:
- Efficient vectorization and cache performance
- Minimal data movement during MPI communication  
- Compatible with Fortran-ordered GSI interfaces
- Support for both single and double precision

# Usage Notes

- Always call `allocate_state!` before using state vector operations
- Use `deallocate_state!` to free memory when done
- State vectors are mutable - operations modify in place
- Thread-safe operations for OpenMP parallelization
- MPI-aware for distributed memory parallelism
"""
mutable struct StateVector{T<:AbstractFloat} <: AbstractStateVector{T}
    # Core 3D meteorological fields
    u::Array{T,3}         # Zonal wind [m/s]
    v::Array{T,3}         # Meridional wind [m/s] 
    t::Array{T,3}         # Temperature [K]
    q::Array{T,3}         # Specific humidity [kg/kg]
    prse::Array{T,3}      # Pressure on edges [Pa]
    oz::Array{T,3}        # Ozone mixing ratio [kg/kg]
    
    # Cloud variables (3D)
    cw::Array{T,3}        # Cloud water [kg/kg]
    ql::Array{T,3}        # Liquid water content [kg/kg]
    qi::Array{T,3}        # Ice water content [kg/kg] 
    qr::Array{T,3}        # Rain water content [kg/kg]
    qs::Array{T,3}        # Snow content [kg/kg]
    qg::Array{T,3}        # Graupel content [kg/kg]
    
    # Surface variables (2D)
    ps::Matrix{T}         # Surface pressure [Pa]
    sst::Matrix{T}        # Sea surface temperature [K]
    stl::Matrix{T}        # Skin temp over land [K] 
    sti::Matrix{T}        # Skin temp over ice [K]
    
    # Diagnostic/derived variables (3D)
    tv::Array{T,3}        # Virtual temperature [K]
    tsen::Array{T,3}      # Sensible temperature [K]
    w::Array{T,3}         # Vertical velocity [m/s]
    
    # Auxiliary 2D variables  
    gust::Matrix{T}       # Surface wind gust [m/s]
    vis::Matrix{T}        # Visibility [m]
    pblh::Matrix{T}       # PBL height [m]
    wspd10m::Matrix{T}    # 10m wind speed [m/s]
    tcamt::Matrix{T}      # Total cloud amount
    lcbas::Matrix{T}      # Low cloud base [m]
    
    # Chemical species (3D) - optional
    co::Array{T,3}        # Carbon monoxide [kg/kg]
    co2::Array{T,3}       # Carbon dioxide [kg/kg]
    
    # Configuration and metadata
    config::AbstractAnalysisConfig
    variables::AnalysisVariables
    grid::GridDimensions
    is_allocated::Bool
    memory_usage::Dict{String,Int}
    
    # MPI information
    mpi_rank::Int
    mpi_size::Int
    local_grid_start::Tuple{Int,Int,Int}
    local_grid_end::Tuple{Int,Int,Int}
    
    function StateVector{T}(config::AbstractAnalysisConfig,
                           variables::AnalysisVariables = default_analysis_variables()) where T
        
        # Initialize MPI info
        mpi_rank = MPI.Initialized() ? MPI.Comm_rank(MPI.COMM_WORLD) : 0
        mpi_size = MPI.Initialized() ? MPI.Comm_size(MPI.COMM_WORLD) : 1
        
        grid = GridDimensions(config.grid_size)
        
        # Calculate local domain decomposition (simplified)
        ny_local = grid.ny ÷ mpi_size
        j_start = mpi_rank * ny_local + 1
        j_end = min((mpi_rank + 1) * ny_local, grid.ny)
        
        sv = new{T}(
            # Initialize empty arrays - will be allocated later
            Array{T,3}(undef, 0, 0, 0),  # u
            Array{T,3}(undef, 0, 0, 0),  # v
            Array{T,3}(undef, 0, 0, 0),  # t
            Array{T,3}(undef, 0, 0, 0),  # q
            Array{T,3}(undef, 0, 0, 0),  # prse
            Array{T,3}(undef, 0, 0, 0),  # oz
            Array{T,3}(undef, 0, 0, 0),  # cw
            Array{T,3}(undef, 0, 0, 0),  # ql
            Array{T,3}(undef, 0, 0, 0),  # qi
            Array{T,3}(undef, 0, 0, 0),  # qr
            Array{T,3}(undef, 0, 0, 0),  # qs
            Array{T,3}(undef, 0, 0, 0),  # qg
            Matrix{T}(undef, 0, 0),      # ps
            Matrix{T}(undef, 0, 0),      # sst
            Matrix{T}(undef, 0, 0),      # stl
            Matrix{T}(undef, 0, 0),      # sti
            Array{T,3}(undef, 0, 0, 0),  # tv
            Array{T,3}(undef, 0, 0, 0),  # tsen
            Array{T,3}(undef, 0, 0, 0),  # w
            Matrix{T}(undef, 0, 0),      # gust
            Matrix{T}(undef, 0, 0),      # vis
            Matrix{T}(undef, 0, 0),      # pblh
            Matrix{T}(undef, 0, 0),      # wspd10m
            Matrix{T}(undef, 0, 0),      # tcamt
            Matrix{T}(undef, 0, 0),      # lcbas
            Array{T,3}(undef, 0, 0, 0),  # co
            Array{T,3}(undef, 0, 0, 0),  # co2
            config,                      # config
            variables,                   # variables
            grid,                        # grid
            false,                       # is_allocated
            Dict{String,Int}(),          # memory_usage
            mpi_rank,                    # mpi_rank
            mpi_size,                    # mpi_size
            (1, j_start, 1),            # local_grid_start
            (grid.nx, j_end, grid.nz)   # local_grid_end
        )
        
        # Allocate memory
        allocate_state!(sv)
        
        return sv
    end
end

# Convenience constructor
StateVector(config::AbstractAnalysisConfig, args...; kwargs...) = 
    StateVector{config.precision}(config, args...; kwargs...)

# ============================================================================
# Memory Management Functions
# ============================================================================

"""
    allocate_state!(state::StateVector{T}) where T

Allocate memory for all state vector variables based on local grid dimensions.

This function allocates memory for all physical variables in the state vector,
taking into account MPI domain decomposition. Memory usage is tracked and 
statistics are updated for monitoring purposes.

# Arguments
- `state::StateVector{T}`: State vector to allocate memory for

# Returns
- Modified state vector with allocated memory

# Mathematical Details
The memory layout follows the GSI convention:
- 3D fields: (nx, ny_local, nz) where ny_local = ny/npe for domain decomposition  
- 2D fields: (nx, ny_local)
- Total memory: Σ(field_size * sizeof(T)) for all fields

# Implementation Notes
- Respects MPI domain decomposition boundaries
- Tracks memory usage per variable type
- Initializes arrays with zeros for numerical stability
- Thread-safe allocation for OpenMP compatibility

# Example
```julia
state = StateVector(config)  # Allocation happens automatically
# Or manually:
# allocate_state!(state)
```
"""
function allocate_state!(state::StateVector{T}) where T
    if state.is_allocated
        return state
    end
    
    # Local grid dimensions for this MPI rank
    nx = state.grid.nx
    ny_local = state.local_grid_end[2] - state.local_grid_start[2] + 1
    nz = state.grid.nz
    
    # Track memory usage
    bytes_per_element = sizeof(T)
    total_memory = 0
    
    # Allocate 3D meteorological fields
    field_3d_size = nx * ny_local * nz
    field_3d_memory = field_3d_size * bytes_per_element
    
    state.u = zeros(T, nx, ny_local, nz)
    state.v = zeros(T, nx, ny_local, nz)
    state.t = zeros(T, nx, ny_local, nz)
    state.q = zeros(T, nx, ny_local, nz)
    state.prse = zeros(T, nx, ny_local, nz)
    state.oz = zeros(T, nx, ny_local, nz)
    
    state.memory_usage["core_3d"] = 6 * field_3d_memory
    total_memory += 6 * field_3d_memory
    
    # Allocate cloud variables (3D)
    state.cw = zeros(T, nx, ny_local, nz)
    state.ql = zeros(T, nx, ny_local, nz)
    state.qi = zeros(T, nx, ny_local, nz)
    state.qr = zeros(T, nx, ny_local, nz)
    state.qs = zeros(T, nx, ny_local, nz)
    state.qg = zeros(T, nx, ny_local, nz)
    
    state.memory_usage["cloud_3d"] = 6 * field_3d_memory
    total_memory += 6 * field_3d_memory
    
    # Allocate diagnostic/derived variables (3D)
    state.tv = zeros(T, nx, ny_local, nz)
    state.tsen = zeros(T, nx, ny_local, nz)
    state.w = zeros(T, nx, ny_local, nz)
    
    state.memory_usage["diagnostic_3d"] = 3 * field_3d_memory
    total_memory += 3 * field_3d_memory
    
    # Allocate 2D surface variables
    field_2d_size = nx * ny_local
    field_2d_memory = field_2d_size * bytes_per_element
    
    state.ps = zeros(T, nx, ny_local)
    state.sst = zeros(T, nx, ny_local)
    state.stl = zeros(T, nx, ny_local)
    state.sti = zeros(T, nx, ny_local)
    
    state.memory_usage["surface_2d"] = 4 * field_2d_memory
    total_memory += 4 * field_2d_memory
    
    # Allocate auxiliary 2D variables
    state.gust = zeros(T, nx, ny_local)
    state.vis = zeros(T, nx, ny_local)
    state.pblh = zeros(T, nx, ny_local)
    state.wspd10m = zeros(T, nx, ny_local)
    state.tcamt = zeros(T, nx, ny_local)
    state.lcbas = zeros(T, nx, ny_local)
    
    state.memory_usage["auxiliary_2d"] = 6 * field_2d_memory
    total_memory += 6 * field_2d_memory
    
    # Allocate chemical species (3D) - optional
    state.co = zeros(T, nx, ny_local, nz)
    state.co2 = zeros(T, nx, ny_local, nz)
    
    state.memory_usage["chemistry_3d"] = 2 * field_3d_memory
    total_memory += 2 * field_3d_memory
    
    state.memory_usage["total"] = total_memory
    state.is_allocated = true
    
    # Log allocation success for rank 0
    if state.mpi_rank == 0
        println("StateVector allocated: $(total_memory ÷ 1024^2) MB total memory")
        println("  Grid dimensions: $(nx) x $(ny_local) x $(nz) (local)")
        println("  3D fields: $(field_3d_size) elements each")
        println("  2D fields: $(field_2d_size) elements each")
    end
    
    return state
end

"""
    deallocate_state!(state::StateVector)

Deallocate memory for state vector variables and reset allocation status.

# Arguments
- `state::StateVector`: State vector to deallocate

# Returns
- Modified state vector with deallocated memory
"""
function deallocate_state!(state::StateVector)
    if !state.is_allocated
        return state
    end
    
    # Resize all arrays to minimal size to free memory
    # 3D fields
    state.u = Array{eltype(state.u),3}(undef, 0, 0, 0)
    state.v = Array{eltype(state.v),3}(undef, 0, 0, 0)
    state.t = Array{eltype(state.t),3}(undef, 0, 0, 0)
    state.q = Array{eltype(state.q),3}(undef, 0, 0, 0)
    state.prse = Array{eltype(state.prse),3}(undef, 0, 0, 0)
    state.oz = Array{eltype(state.oz),3}(undef, 0, 0, 0)
    state.cw = Array{eltype(state.cw),3}(undef, 0, 0, 0)
    state.ql = Array{eltype(state.ql),3}(undef, 0, 0, 0)
    state.qi = Array{eltype(state.qi),3}(undef, 0, 0, 0)
    state.qr = Array{eltype(state.qr),3}(undef, 0, 0, 0)
    state.qs = Array{eltype(state.qs),3}(undef, 0, 0, 0)
    state.qg = Array{eltype(state.qg),3}(undef, 0, 0, 0)
    state.tv = Array{eltype(state.tv),3}(undef, 0, 0, 0)
    state.tsen = Array{eltype(state.tsen),3}(undef, 0, 0, 0)
    state.w = Array{eltype(state.w),3}(undef, 0, 0, 0)
    state.co = Array{eltype(state.co),3}(undef, 0, 0, 0)
    state.co2 = Array{eltype(state.co2),3}(undef, 0, 0, 0)
    
    # 2D fields
    state.ps = Matrix{eltype(state.ps)}(undef, 0, 0)
    state.sst = Matrix{eltype(state.sst)}(undef, 0, 0)
    state.stl = Matrix{eltype(state.stl)}(undef, 0, 0)
    state.sti = Matrix{eltype(state.sti)}(undef, 0, 0)
    state.gust = Matrix{eltype(state.gust)}(undef, 0, 0)
    state.vis = Matrix{eltype(state.vis)}(undef, 0, 0)
    state.pblh = Matrix{eltype(state.pblh)}(undef, 0, 0)
    state.wspd10m = Matrix{eltype(state.wspd10m)}(undef, 0, 0)
    state.tcamt = Matrix{eltype(state.tcamt)}(undef, 0, 0)
    state.lcbas = Matrix{eltype(state.lcbas)}(undef, 0, 0)
    
    # Clear memory usage tracking
    empty!(state.memory_usage)
    state.is_allocated = false
    
    # Force garbage collection
    GC.gc()
    
    return state
end

# ============================================================================
# State Vector Operations
# ============================================================================

"""
    compute_state_norms(state::StateVector{T}; variable::Union{String,Nothing}=nothing) where T

Compute statistical norms (min, max, mean, sum) for state vector variables.

This function computes comprehensive statistics for all or specific variables in the
state vector, with MPI reduction across all processes to provide global statistics.
Based on GSI's `norms_vars` subroutine.

# Arguments
- `state::StateVector{T}`: Input state vector
- `variable::Union{String,Nothing}=nothing`: Specific variable name (e.g., "t", "q") or nothing for all variables

# Returns
- `Dict{String, NamedTuple}`: Dictionary mapping variable names to statistics
  Each entry contains: `(min=..., max=..., mean=..., sum=..., count=...)`

# Mathematical Details
For each variable field `x(i,j,k)`:
- `min = min(x)`: Global minimum value
- `max = max(x)`: Global maximum value  
- `sum = Σᵢⱼₖ x(i,j,k)`: Global sum across all grid points
- `count = N`: Total number of grid points
- `mean = sum/count`: Global mean value

# MPI Reduction
Statistics are computed locally on each process then reduced using:
- `MPI.Allreduce` with `MPI.MIN` for minimum values
- `MPI.Allreduce` with `MPI.MAX` for maximum values  
- `MPI.Allreduce` with `MPI.SUM` for sums and counts

# Example
```julia
norms = compute_state_norms(state)
println("Temperature: min=$(norms["t"].min), max=$(norms["t"].max)")

# Or for specific variable
t_norms = compute_state_norms(state, variable="t")
```
"""
function compute_state_norms(state::StateVector{T}; variable::Union{String,Nothing}=nothing) where T
    @assert state.is_allocated "State vector must be allocated"
    
    norms = Dict{String, NamedTuple}()
    
    # Define variables to process
    variables_to_process = if variable !== nothing
        [variable]
    else
        ["u", "v", "t", "q", "prse", "oz", "cw", "ql", "qi", "qr", "qs", "qg",
         "ps", "sst", "stl", "sti", "tv", "tsen", "w", "gust", "vis", "pblh",
         "wspd10m", "tcamt", "lcbas", "co", "co2"]
    end
    
    for var_name in variables_to_process
        if hasfield(StateVector{T}, Symbol(var_name))
            field = getfield(state, Symbol(var_name))
            
            # Skip if field is empty (not allocated or optional)
            if isempty(field)
                continue
            end
            
            # Local statistics
            local_min = minimum(field)
            local_max = maximum(field)
            local_sum = sum(field)
            local_count = length(field)
            
            # MPI reduction (if MPI is initialized)
            if MPI.Initialized() && state.mpi_size > 1
                global_min = MPI.Allreduce(local_min, MPI.MIN, MPI.COMM_WORLD)
                global_max = MPI.Allreduce(local_max, MPI.MAX, MPI.COMM_WORLD)
                global_sum = MPI.Allreduce(local_sum, MPI.SUM, MPI.COMM_WORLD)
                global_count = MPI.Allreduce(local_count, MPI.SUM, MPI.COMM_WORLD)
            else
                global_min = local_min
                global_max = local_max
                global_sum = local_sum
                global_count = local_count
            end
            
            global_mean = global_sum / global_count
            
            norms[var_name] = (
                min = global_min,
                max = global_max,
                mean = global_mean,
                sum = global_sum,
                count = global_count
            )
        end
    end
    
    return norms
end

"""
    dot_product_state(state1::StateVector{T}, state2::StateVector{T}; 
                     variable::Union{String,Nothing}=nothing) where T

Compute dot product between two state vectors or specific variables.

Based on GSI's `dot_prod_st` function with MPI reduction support.

# Arguments
- `state1::StateVector{T}`: First state vector
- `state2::StateVector{T}`: Second state vector
- `variable::Union{String,Nothing}=nothing`: Specific variable name or nothing for all variables

# Returns
- `T`: Dot product value (scalar)

# Mathematical Details
For vectors `x` and `y`, computes: `⟨x,y⟩ = Σᵢ xᵢyᵢ`

If `variable` is specified, computes dot product only for that variable:
`⟨x[var],y[var]⟩ = Σᵢⱼₖ x[var](i,j,k) * y[var](i,j,k)`

If `variable=nothing`, computes total dot product across all variables:
`⟨x,y⟩ = Σᵥₐᵣ ⟨x[var],y[var]⟩`

# Example
```julia
# Full state vector dot product
result = dot_product_state(state1, state2)

# Temperature field only
temp_dot = dot_product_state(state1, state2, variable="t")
```
"""
function dot_product_state(state1::StateVector{T}, state2::StateVector{T}; 
                          variable::Union{String,Nothing}=nothing) where T
    @assert state1.is_allocated && state2.is_allocated "State vectors must be allocated"
    @assert size(state1.u) == size(state2.u) "State vectors must have compatible dimensions"
    
    local_result = zero(T)
    
    # Define variables to process
    variables_to_process = if variable !== nothing
        [variable]
    else
        ["u", "v", "t", "q", "prse", "oz", "cw", "ql", "qi", "qr", "qs", "qg",
         "ps", "sst", "stl", "sti", "tv", "tsen", "w", "gust", "vis", "pblh",
         "wspd10m", "tcamt", "lcbas", "co", "co2"]
    end
    
    for var_name in variables_to_process
        if hasfield(StateVector{T}, Symbol(var_name))
            field1 = getfield(state1, Symbol(var_name))
            field2 = getfield(state2, Symbol(var_name))
            
            # Skip if fields are empty or incompatible
            if isempty(field1) || isempty(field2) || size(field1) != size(field2)
                continue
            end
            
            # Compute local contribution
            local_result += dot(field1, field2)
        end
    end
    
    # MPI reduction (if MPI is initialized)
    if MPI.Initialized() && state1.mpi_size > 1
        global_result = MPI.Allreduce(local_result, MPI.SUM, MPI.COMM_WORLD)
    else
        global_result = local_result
    end
    
    return global_result
end

"""
    set_random_state!(state::StateVector{T}; preserve_physics::Bool=true) where T

Initialize state vector with random values while preserving physical relationships.

Based on GSI's `set_random_st` subroutine with physical consistency constraints.

# Arguments
- `state::StateVector{T}`: State vector to initialize
- `preserve_physics::Bool=true`: Whether to enforce physical consistency

# Mathematical Details
Random initialization follows:
1. Generate independent random fields for each variable: `x ~ U(0,1)`
2. If `preserve_physics=true`, apply constraints:
   - Pressure monotonically decreases with height
   - Temperature and humidity satisfy saturation limits
   - Virtual and sensible temperature consistency: `Tv = T(1 + 0.61q)`

# Physical Constraints Applied
- Surface pressure: `ps > 50000 Pa` (minimum realistic value)
- Temperature: `150K < T < 350K` (reasonable atmospheric range)
- Specific humidity: `0 ≤ q ≤ qsat(T,p)` (saturation limit)
- Mixing ratios: All ≥ 0 (positive definite)

# Example
```julia
# Random initialization with physics
set_random_state!(state)

# Pure random (no physical constraints)
set_random_state!(state, preserve_physics=false)
```
"""
function set_random_state!(state::StateVector{T}; preserve_physics::Bool=true) where T
    @assert state.is_allocated "State vector must be allocated"
    
    # Import Random for seed control (would be at module level in real implementation)
    import Random
    
    # Set random seed based on MPI rank for reproducibility
    Random.seed!(42 + state.mpi_rank)
    
    # Initialize all fields with random values [0,1]
    rand!(state.u)
    rand!(state.v)
    rand!(state.t)
    rand!(state.q)
    rand!(state.prse)
    rand!(state.oz)
    rand!(state.cw)
    rand!(state.ql)
    rand!(state.qi)
    rand!(state.qr)
    rand!(state.qs)
    rand!(state.qg)
    rand!(state.ps)
    rand!(state.sst)
    rand!(state.stl)
    rand!(state.sti)
    rand!(state.tv)
    rand!(state.tsen)
    rand!(state.w)
    rand!(state.gust)
    rand!(state.vis)
    rand!(state.pblh)
    rand!(state.wspd10m)
    rand!(state.tcamt)
    rand!(state.lcbas)
    rand!(state.co)
    rand!(state.co2)
    
    if preserve_physics
        # Apply physical scaling and constraints
        
        # Scale winds to realistic range [-50, 50] m/s
        state.u .= (state.u .- 0.5) .* 100
        state.v .= (state.v .- 0.5) .* 100
        
        # Scale temperature to [200, 320] K
        state.t .= 200 .+ state.t .* 120
        
        # Scale specific humidity to [0, 0.02] kg/kg
        state.q .*= 0.02
        
        # Scale surface pressure to [70000, 105000] Pa
        state.ps .= 70000 .+ state.ps .* 35000
        
        # Ensure pressure decreases with height (simplified)
        for k in 1:size(state.prse, 3)
            state.prse[:, :, k] .= state.ps .* (1.0 - 0.1 * (k-1) / size(state.prse, 3))
        end
        
        # Scale temperature fields
        state.sst .= 270 .+ state.sst .* 30  # [270, 300] K
        state.stl .= 250 .+ state.stl .* 50  # [250, 300] K  
        state.sti .= 220 .+ state.sti .* 30  # [220, 250] K
        
        # Virtual temperature consistency: tv = t * (1 + 0.61*q)
        state.tv .= state.t .* (1 .+ 0.61 .* state.q)
        state.tsen .= state.t  # Assume dry air for simplicity
        
        # Scale mixing ratios to small positive values
        state.oz .*= 1e-5   # Ozone [0, 1e-5] kg/kg
        state.cw .*= 1e-3   # Cloud water [0, 1e-3] kg/kg
        state.ql .*= 1e-3   # Liquid water
        state.qi .*= 1e-3   # Ice water
        state.qr .*= 1e-4   # Rain water
        state.qs .*= 1e-4   # Snow
        state.qg .*= 1e-5   # Graupel
        state.co .*= 1e-6   # CO mixing ratio
        state.co2 .*= 4e-4  # CO2 mixing ratio (~400 ppm)
        
        # Scale auxiliary variables
        state.w .= (state.w .- 0.5) .* 10      # Vertical velocity [-5, 5] m/s
        state.gust .*= 30                      # Wind gust [0, 30] m/s
        state.vis .= 1000 .+ state.vis .* 9000 # Visibility [1, 10] km
        state.pblh .*= 3000                    # PBL height [0, 3000] m
        state.wspd10m .*= 25                   # Wind speed [0, 25] m/s
        state.tcamt .*= 1.0                    # Cloud fraction [0, 1]
        state.lcbas .*= 5000                   # Cloud base [0, 5000] m
    end
    
    return state
end

# ============================================================================
# Control-to-State Transformations (from control2state.f90)
# ============================================================================

"""
    control_to_state!(control_vector, state::StateVector{T}, config::AbstractAnalysisConfig) where T

Transform control variables to physical state variables.

This function performs the critical transformation from control variables 
(in transformed/normalized space) to physical state variables, implementing
the core functionality of GSI's `control2state.f90`.

# Arguments
- `control_vector`: Control vector containing transformed variables
- `state::StateVector{T}`: Output state vector (modified in place)
- `config::AbstractAnalysisConfig`: Analysis configuration

# Transformations Performed
1. **Streamfunction/Velocity Potential → Wind Components**
   - `u = ∂ψ/∂y - ∂φ/∂x` (zonal wind)
   - `v = -∂ψ/∂x - ∂φ/∂y` (meridional wind)

2. **Normalized Humidity → Specific Humidity** 
   - `q = qnorm * qsat(T,p)` where qsat is saturation mixing ratio

3. **Virtual Temperature → Sensible Temperature**
   - `T = Tv / (1 + 0.61*q)` (dry temperature)

4. **Pressure Coordinate Transformations**
   - `prse = f(ps, T)` using hydrostatic balance

5. **Cloud Microphysics** (if enabled)
   - Transform cloud water to hydrometeor species (ql, qi, qr, qs, qg)

# Mathematical Framework
The transformation implements: `x_physical = T(x_control)`
where T represents the nonlinear transformation operator.

# Example
```julia
control_to_state!(control_vec, state, config)
# state now contains physical variables ready for observation operators
```
"""
function control_to_state!(control_vector, state::StateVector{T}, config::AbstractAnalysisConfig) where T
    @assert state.is_allocated "State vector must be allocated"
    
    # This is a simplified implementation - the full version would include:
    # 1. Streamfunction/velocity potential to u,v transformation
    # 2. Humidity normalization and conversion
    # 3. Pressure coordinate transformations  
    # 4. Balance constraints
    # 5. Cloud microphysics transformations
    
    if state.mpi_rank == 0
        println("Applying control-to-state transformations...")
    end
    
    # Placeholder: Direct copy from control vector to state
    # In real implementation, this would involve complex transformations
    streamfunction_to_winds!(state, config)
    pressure_transforms!(state, config)
    humidity_transforms!(state, config)
    balance_constraints!(state, config)
    
    if config.use_hybrid
        cloud_microphysics_transforms!(state, config)
    end
    
    return state
end

"""
    streamfunction_to_winds!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Transform streamfunction and velocity potential to wind components.

# Mathematical Details
Given streamfunction ψ and velocity potential φ:
- `u = ∂ψ/∂y - ∂φ/∂x` (zonal wind)
- `v = -∂ψ/∂x - ∂φ/∂y` (meridional wind)

This transformation ensures the wind field satisfies both rotational and
divergent components of the flow.
"""
function streamfunction_to_winds!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Placeholder implementation
    # Real implementation would use spectral derivatives or finite differences
    # to compute ∂ψ/∂x, ∂ψ/∂y, ∂φ/∂x, ∂φ/∂y
    
    # For now, assume u,v are already computed or use simple approximation
    # This would involve FFT transforms for spectral derivatives
    
    return state
end

"""
    pressure_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Apply pressure coordinate transformations and hydrostatic balance.

Computes pressure on layer edges using hydrostatic equation:
`dp/dz = -ρg = -p/(RT) * g`
"""
function pressure_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Simplified hydrostatic balance
    # Real implementation would use GSI's getprs_tl routine
    
    g = T(9.81)  # gravitational acceleration
    R = T(287.0) # gas constant for dry air
    
    # Compute pressure at each level (simplified)
    for k in 1:size(state.prse, 3)
        # Simple exponential decrease with height
        height_factor = T(1.0 - 0.1 * (k-1) / size(state.prse, 3))
        state.prse[:, :, k] = state.ps .* height_factor
    end
    
    return state
end

"""
    humidity_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Transform normalized relative humidity to specific humidity.

# Mathematical Details  
`q = qnorm * qsat(T,p)` where:
- `qnorm`: Normalized humidity from control vector [0,1]
- `qsat(T,p)`: Saturation mixing ratio as function of temperature and pressure
- `q`: Specific humidity [kg/kg]
"""
function humidity_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Simplified saturation vapor pressure calculation
    # Real implementation would use GSI's normal_rh_to_q routine
    
    # Ensure humidity is within reasonable bounds
    state.q = clamp.(state.q, T(0.0), T(0.025))  # Max ~25 g/kg
    
    # Virtual temperature consistency: tv = t * (1 + 0.61*q)
    state.tv .= state.t .* (T(1.0) .+ T(0.61) .* state.q)
    
    # Sensible temperature (assuming minimal correction)
    state.tsen .= state.t
    
    return state
end

"""
    balance_constraints!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Apply balance constraints (geostrophic, hydrostatic) to ensure physical consistency.

Balance constraints include:
1. **Geostrophic Balance**: Wind-pressure relationship at large scales
2. **Hydrostatic Balance**: Vertical pressure gradient balances gravity
3. **Thermal Wind**: Relationship between temperature gradient and wind shear
"""
function balance_constraints!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Placeholder for balance constraints
    # Real implementation would use GSI's balance routines
    
    # Ensure temperature is within reasonable bounds
    state.t = clamp.(state.t, T(180.0), T(350.0))  # 180K to 350K
    
    # Ensure surface pressure is positive and reasonable
    state.ps = clamp.(state.ps, T(50000.0), T(110000.0))  # 500 to 1100 hPa
    
    return state
end

"""
    cloud_microphysics_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Transform cloud water to individual hydrometeor species.

Based on GSI's cloud microphysics schemes, converts total cloud water
to specific hydrometeor categories:
- `ql`: Liquid water content
- `qi`: Ice water content  
- `qr`: Rain water content
- `qs`: Snow content
- `qg`: Graupel content
"""
function cloud_microphysics_transforms!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Simplified cloud partitioning
    # Real implementation would use cw2hydro_tl from GSI
    
    # Simple temperature-based partitioning
    T_freeze = T(273.15)  # Freezing point
    
    for k in 1:size(state.cw, 3)
        for j in 1:size(state.cw, 2)
            for i in 1:size(state.cw, 1)
                cw_total = state.cw[i, j, k]
                temp = state.t[i, j, k]
                
                if temp > T_freeze
                    # Warm cloud: mostly liquid
                    state.ql[i, j, k] = T(0.8) * cw_total
                    state.qi[i, j, k] = T(0.0) * cw_total
                    state.qr[i, j, k] = T(0.2) * cw_total
                    state.qs[i, j, k] = T(0.0) * cw_total
                    state.qg[i, j, k] = T(0.0) * cw_total
                else
                    # Cold cloud: mostly ice
                    state.ql[i, j, k] = T(0.1) * cw_total
                    state.qi[i, j, k] = T(0.7) * cw_total
                    state.qr[i, j, k] = T(0.0) * cw_total
                    state.qs[i, j, k] = T(0.2) * cw_total
                    state.qg[i, j, k] = T(0.0) * cw_total
                end
            end
        end
    end
    
    return state
end

# ============================================================================
# State-to-Model Transformations (from control2model.f90)
# ============================================================================

"""
    state_to_model!(state::StateVector{T}, model_state, config::AbstractAnalysisConfig) where T

Prepare state variables for model integration by applying sqrt(B) transformation.

This function implements the functionality of GSI's `control2model.f90`,
applying the square root of the background error covariance matrix to
convert from control space to model space.

# Arguments
- `state::StateVector{T}`: Input state vector
- `model_state`: Output model state (modified in place)
- `config::AbstractAnalysisConfig`: Analysis configuration

# Mathematical Framework
Applies the transformation: `x_model = B^(1/2) * x_control`
where B^(1/2) is the square root of background error covariance.

# Transformations Include
1. **Background Error Covariance**: Apply sqrt(B) via spectral transforms
2. **Balance Constraints**: Ensure geostrophic and hydrostatic balance
3. **Physical Consistency**: Maintain thermodynamic relationships
4. **Boundary Conditions**: Apply appropriate boundary constraints

# Example
```julia
state_to_model!(state, model_vars, config)
# model_vars now ready for forecast model integration
```
"""
function state_to_model!(state::StateVector{T}, model_state, config::AbstractAnalysisConfig) where T
    @assert state.is_allocated "State vector must be allocated"
    
    if state.mpi_rank == 0
        println("Applying state-to-model transformations...")
    end
    
    # Apply sqrt(B) transformation (simplified)
    # Real implementation would use GSI's ckgcov routine
    apply_background_covariance!(state, config)
    
    # Apply additional balance constraints for model consistency
    enforce_model_constraints!(state, config)
    
    # Copy to model state format (implementation-dependent)
    # This would depend on the specific model interface
    
    return state
end

"""
    apply_background_covariance!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Apply square root of background error covariance matrix.

This implements the core B^(1/2) transformation using recursive filters,
spectral transforms, or other methods depending on the covariance model.
"""
function apply_background_covariance!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Placeholder for sqrt(B) application
    # Real implementation would use spectral transforms and recursive filters
    
    # Simple scaling as placeholder (would be much more complex in reality)
    scaling_factor = T(0.1)  # Typical analysis increment magnitude
    
    state.u .*= scaling_factor
    state.v .*= scaling_factor  
    state.t .*= scaling_factor
    state.q .*= scaling_factor
    state.ps .*= scaling_factor
    
    return state
end

"""
    enforce_model_constraints!(state::StateVector{T}, config::AbstractAnalysisConfig) where T

Enforce physical and numerical constraints required by the forecast model.

These constraints ensure the analysis increment is compatible with
the model's numerical schemes and physical parameterizations.
"""
function enforce_model_constraints!(state::StateVector{T}, config::AbstractAnalysisConfig) where T
    # Ensure all mixing ratios are non-negative
    state.q = max.(state.q, T(0.0))
    state.oz = max.(state.oz, T(0.0))
    state.cw = max.(state.cw, T(0.0))
    state.ql = max.(state.ql, T(0.0))
    state.qi = max.(state.qi, T(0.0))
    state.qr = max.(state.qr, T(0.0))
    state.qs = max.(state.qs, T(0.0))
    state.qg = max.(state.qg, T(0.0))
    
    # Ensure surface pressure is positive
    state.ps = max.(state.ps, T(10000.0))  # Minimum 100 hPa
    
    # Ensure reasonable temperature bounds
    state.t = clamp.(state.t, T(100.0), T(400.0))
    state.tv = clamp.(state.tv, T(100.0), T(400.0))
    state.tsen = clamp.(state.tsen, T(100.0), T(400.0))
    
    return state
end

# ============================================================================
# Utility Functions
# ============================================================================

"""
    inquire_state_memory(state::StateVector{T}) where T

Report memory usage statistics for the state vector.

Based on GSI's `inquire_state` subroutine.
"""
function inquire_state_memory(state::StateVector{T}) where T
    if !state.is_allocated
        println("State vector is not allocated")
        return
    end
    
    total_mb = state.memory_usage["total"] ÷ 1024^2
    
    println("State Vector Memory Usage:")
    println("  Grid dimensions: $(state.grid.nx) × $(state.local_grid_end[2] - state.local_grid_start[2] + 1) × $(state.grid.nz) (local)")
    println("  MPI rank: $(state.mpi_rank) of $(state.mpi_size)")
    println("  Total memory: $(total_mb) MB")
    
    for (category, bytes) in state.memory_usage
        if category != "total"
            mb = bytes ÷ 1024^2
            println("    $(category): $(mb) MB")
        end
    end
end

"""
    print_state_norms(state::StateVector{T}; variable::Union{String,Nothing}=nothing) where T

Print formatted norm statistics for state vector variables.

Based on GSI's `prt_state_norms` subroutines.
"""
function print_state_norms(state::StateVector{T}; variable::Union{String,Nothing}=nothing) where T
    norms = compute_state_norms(state, variable=variable)
    
    if state.mpi_rank == 0
        println("State Vector Norms:")
        println("Variable    Mean              Min               Max")
        println("--------    ----              ---               ---")
        
        for (var_name, stats) in sort(collect(norms))
            @printf("%-8s    %12.6e    %12.6e    %12.6e\n", 
                    var_name, stats.mean, stats.min, stats.max)
        end
    end
end

end # module StateVectors